#include "MessagePool.h"
#include <algorithm>

namespace SeanNetwork {

// TCP重传策略实现
TcpResendStrategy::TcpResendStrategy(uint32_t initialInterval, uint32_t maxInterval, uint32_t maxResendCount)
    : m_InitialInterval(initialInterval), m_MaxInterval(maxInterval), m_MaxResendCount(maxResendCount) {
}

uint32_t TcpResendStrategy::GetNextResendInterval(uint32_t resendCount) {
    /* 功能：计算TCP重传间隔（指数退避）
       参数：resendCount-当前重传次数
       返回：下次重传间隔（毫秒） */
    uint32_t interval = m_InitialInterval;
    for (uint32_t i = 0; i < resendCount && interval < m_MaxInterval / 2; ++i) {
        interval *= 2;  // 指数退避
    }
    return std::min(interval, m_MaxInterval);
}

uint32_t TcpResendStrategy::GetMaxResendCount() {
    return m_MaxResendCount;
}

// KCP重传策略实现
KcpResendStrategy::KcpResendStrategy(uint32_t fixedInterval, uint32_t maxResendCount)
    : m_FixedInterval(fixedInterval), m_MaxResendCount(maxResendCount) {
}

uint32_t KcpResendStrategy::GetNextResendInterval(uint32_t resendCount) {
    /* 功能：计算KCP重传间隔（固定间隔，快速重传）
       参数：resendCount-当前重传次数
       返回：固定重传间隔（毫秒） */
    return m_FixedInterval;
}

uint32_t KcpResendStrategy::GetMaxResendCount() {
    return m_MaxResendCount;
}

// 消息池实现
MessagePool::MessagePool() : m_MessageTtl(5000) {
    // 设置默认重传策略
    m_ResendStrategies[ChannelType::TCP] = std::make_shared<TcpResendStrategy>();
    m_ResendStrategies[ChannelType::KCP] = std::make_shared<KcpResendStrategy>();
    
    // 初始化优先级队列
    for (int i = 0; i <= 3; ++i) {
        m_PriorityQueues[i] = std::queue<uint32_t>();
    }
}

MessagePool::~MessagePool() {
    Clear();
}

bool MessagePool::AddMessage(std::shared_ptr<NetworkMessage> message, ChannelType channel) {
    /* 功能：添加消息到发送池
       参数：message-要发送的消息，channel-传输通道
       返回：添加是否成功 */
    if (!message) {
        return false;
    }
    
    std::lock_guard<std::mutex> lock(m_Mutex);
    
    uint32_t messageId = message->GetHeader().MessageId;
    uint8_t priority = message->GetHeader().Priority;
    
    // 如果消息已经在池中，更新它
    auto it = m_PendingMessages.find(messageId);
    if (it != m_PendingMessages.end()) {
        it->second.Message = message;
        it->second.Channel = channel;
        it->second.NextResendTime = CalculateNextResendTime(message, channel);
        return true;
    }
    
    // 添加新消息
    message->UpdateLastSendTime();
    MessageState state(message, channel);
    state.NextResendTime = CalculateNextResendTime(message, channel);
    
    m_PendingMessages[messageId] = state;
    m_PriorityQueues[priority].push(messageId);
    
    return true;
}

bool MessagePool::AcknowledgeMessage(uint32_t messageId, ChannelType channel) {
    /* 功能：确认消息已接收
       参数：messageId-要确认的消息ID，channel-传输通道
       返回：确认是否成功 */
    std::lock_guard<std::mutex> lock(m_Mutex);
    
    auto it = m_PendingMessages.find(messageId);
    if (it != m_PendingMessages.end() && it->second.Channel == channel) {
        // 执行回调
        if (m_AckCallback) {
            m_AckCallback(messageId, true);
        }
        
        // 从池中移除消息
        m_PendingMessages.erase(it);
        return true;
    }
    
    return false;
}

void MessagePool::SetAckCallback(MessageAckCallback callback) {
    std::lock_guard<std::mutex> lock(m_Mutex);
    m_AckCallback = callback;
}

std::vector<std::shared_ptr<NetworkMessage>> MessagePool::CheckTimeoutMessages() {
    /* 功能：检查和处理超时消息
       返回：需要重传的消息列表 */
    std::lock_guard<std::mutex> lock(m_Mutex);
    std::vector<std::shared_ptr<NetworkMessage>> timeoutMessages;
    
    uint64_t currentTime = GetCurrentTimeMillis();
    std::vector<uint32_t> timeoutIds;
    
    // 检查待确认的消息
    for (auto& pair : m_PendingMessages) {
        uint32_t messageId = pair.first;
        auto& state = pair.second;
        auto& message = state.Message;
        
        // 检查消息存活时间
        if (message->GetAliveTime() > static_cast<int64_t>(m_MessageTtl)) {
            timeoutIds.push_back(messageId);
            continue;
        }
        
        // 检查是否需要重传
        if (currentTime >= state.NextResendTime) {
            auto strategy = m_ResendStrategies[state.Channel];
            
            /* 重传触发条件：
               1. 当前时间已超过下次重传时间
               2. 未超过最大重传次数 */
            if (message->GetResendCount() < strategy->GetMaxResendCount()) {
                message->IncrementResendCount();
                message->UpdateLastSendTime();
                state.NextResendTime = CalculateNextResendTime(message, state.Channel);
                timeoutMessages.push_back(message);
            } else {
                // 超过最大重传次数，标记为超时
                timeoutIds.push_back(messageId);
            }
        }
    }
    
    // 处理超时消息
    for (uint32_t messageId : timeoutIds) {
        auto it = m_PendingMessages.find(messageId);
        if (it != m_PendingMessages.end()) {
            HandleMessageTimeout(messageId, it->second.Channel);
        }
    }
    
    return timeoutMessages;
}

void MessagePool::SetMessageTimeToLive(uint64_t ttlMilliseconds) {
    std::lock_guard<std::mutex> lock(m_Mutex);
    m_MessageTtl = ttlMilliseconds;
}

void MessagePool::SetResendStrategy(ChannelType channel, std::shared_ptr<IResendStrategy> strategy) {
    if (!strategy) {
        return;
    }
    
    std::lock_guard<std::mutex> lock(m_Mutex);
    m_ResendStrategies[channel] = strategy;
}

std::shared_ptr<NetworkMessage> MessagePool::GetNextPendingMessage() {
    /* 功能：获取待发送的优先消息
       返回：下一个待发送消息，无消息时返回nullptr */
    std::lock_guard<std::mutex> lock(m_Mutex);
    
    // 按优先级从高到低查找待发送消息
    for (int priority = 0; priority <= 3; ++priority) {
        auto& queue = m_PriorityQueues[priority];
        
        while (!queue.empty()) {
            uint32_t messageId = queue.front();
            queue.pop();
            
            auto it = m_PendingMessages.find(messageId);
            if (it != m_PendingMessages.end()) {
                return it->second.Message;
            }
        }
    }
    
    return nullptr;
}

size_t MessagePool::GetPendingMessageCount() const {
    std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(m_Mutex));
    
    size_t count = 0;
    for (const auto& queue : m_PriorityQueues) {
        count += queue.second.size();
    }
    
    return count;
}

size_t MessagePool::GetUnacknowledgedMessageCount(ChannelType channel) const {
    /* 功能：获取指定通道上未确认的消息数
       参数：channel-传输通道
       返回：未确认消息数 */
    std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(m_Mutex));
    
    size_t count = 0;
    for (const auto& pair : m_PendingMessages) {
        if (pair.second.Channel == channel) {
            count++;
        }
    }
    
    return count;
}

void MessagePool::Clear() {
    /* 功能：清空消息池
       说明：释放所有资源，安全地清空消息池 */
    std::lock_guard<std::mutex> lock(m_Mutex);
    
    // 清空优先级队列
    for (auto& queue : m_PriorityQueues) {
        std::queue<uint32_t> empty;
        std::swap(queue.second, empty);
    }
    
    // 清空待确认消息
    m_PendingMessages.clear();
    
    // 清空重传策略映射
    m_ResendStrategies.clear();
}

uint64_t MessagePool::CalculateNextResendTime(const std::shared_ptr<NetworkMessage>& message, ChannelType channel) {
    /* 功能：计算下一次重传时间
       参数：message-需要重传的消息，channel-传输通道
       返回：下次重传的时间点（毫秒时间戳） */
    auto strategy = m_ResendStrategies[channel];
    if (!strategy) {
        return 0;
    }
    
    uint32_t interval = strategy->GetNextResendInterval(message->GetResendCount());
    return message->GetLastSendTime() + interval;
}

void MessagePool::HandleMessageTimeout(uint32_t messageId, ChannelType channel) {
    /* 功能：处理消息超时
       参数：messageId-超时的消息ID，channel-传输通道 */
    // 执行回调
    if (m_AckCallback) {
        m_AckCallback(messageId, false);
    }
    
    // 从池中移除消息
    m_PendingMessages.erase(messageId);
}

bool MessagePool::Initialize() {
    std::lock_guard<std::mutex> lock(m_Mutex);
    
    m_PendingMessages.clear();
    for (auto& queue : m_PriorityQueues) {
        while (!queue.second.empty()) {
            queue.second.pop();
        }
    }
    
    // 默认设置消息TTL为30秒
    m_MessageTtl = 30000;
    
    // 设置默认重传策略
    if (m_ResendStrategies.find(ChannelType::TCP) == m_ResendStrategies.end()) {
        m_ResendStrategies[ChannelType::TCP] = std::make_shared<TcpResendStrategy>();
    }
    
    if (m_ResendStrategies.find(ChannelType::KCP) == m_ResendStrategies.end()) {
        m_ResendStrategies[ChannelType::KCP] = std::make_shared<KcpResendStrategy>();
    }
    
    return true;
}

void MessagePool::Update() {
    // 检查超时消息并执行重传
    CheckTimeoutMessages();
}

void MessagePool::EnqueueOutgoingMessage(std::shared_ptr<NetworkMessage> message) {
    /* 功能：无锁方式添加发送消息
       参数：message-要发送的消息 */
    if (!message) {
        return;
    }
    
    // 直接使用无锁队列入队
    m_OutgoingQueue.Enqueue(message);
}

bool MessagePool::DequeueOutgoingMessage(std::shared_ptr<NetworkMessage>& message) {
    /* 功能：无锁方式获取下一个待发送消息
       参数：message-用于存储获取的消息的引用
       返回：如果有消息则返回true，否则返回false */
    
    // 直接使用无锁队列出队
    return m_OutgoingQueue.Dequeue(message);
}
}